#!/usr/bin/env python

""" good_features_detector.py - Version 0.1 2011-06-08

    Implements the OpenCV Good Features to Track filter on an RGB video stream.

    Created for the Pi Robot Project: http://www.pirobot.org
    Copyright (c) 2011 Patrick Goebel.  All rights reserved.

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.5
    
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details at:
    
    http://www.gnu.org/licenses/gpl.html
      
"""

import roslib
roslib.load_manifest('pi_rein')
from detector import Detector
import rospy
import cv
import sys
from pi_rein.msg import *
from pi_rein.srv import *
from sensor_msgs.msg import Image, PointCloud2
from std_msgs.msg import Header, ColorRGBA
import random
import os

class GoodFeatures(Detector):
    def __init__(self, node_name):
        Detector.__init__(self, node_name)
        
        self.features_pub = rospy.Publisher("detections", DetectionArray)
        
        self.auto_face_tracking = rospy.get_param("~auto_face_tracking", False)
        self.use_depth_for_tracking = rospy.get_param("~use_depth_for_tracking", False)
        self.auto_min_features = rospy.get_param("~auto_min_features", True)
        self.min_features = rospy.get_param("~min_features", 50) # Used only if auto_min_features is False
        self.abs_min_features = rospy.get_param("~abs_min_features", 6)
        self.std_err_xy = rospy.get_param("~std_err_xy", 2.5) 
        self.max_mse = rospy.get_param("~max_mse", 10000)
        self.good_feature_distance = rospy.get_param("~good_feature_distance", 5)
        self.add_feature_distance = rospy.get_param("~add_feature_distance", 10)
        self.feature_type = rospy.get_param("~feature_type", 0) # 0 = Good Features to Track, 1 = SURF
        self.get_surf_also = rospy.get_param("~get_surf_also", False)
        self.expand_scale = 1.1
        
        self.grey = None
        self.small_image = None
        self.pyramid = None
        
        """ Set the Good Features to Track and Lucas-Kanade parameters """
        self.features = []
        self.quality = 0.01
        self.win_size = 10
        self.max_count = 200
        self.flags = 0
        self.frame_count = 0
        
        """ Tracking variables """
        self.track_box = None
        
        rospy.wait_for_message("image", Image, timeout=15)

    def image_callback(self, ros_image):
        (self.image, self.image_size, self.image_width, self.image_height) = self.ros2cv_image(ros_image, "bgr8")
        self.image_stamp = ros_image.header.stamp
        
        try:
            self.detect()
            features = self.getDetections()
            self.features_pub.publish(features)
            if self.use_display and self.show_features:
                self.draw_markers(features.header, features.detections)
        except:
            pass
                    
    def detect(self):
        """ Clear any left over detections """
        self.detections = DetectionArray()
        
        """ Initialize intermediate images if necessary """
        if not self.pyramid:
            self.grey = cv.CreateImage(self.image_size, 8, 1)
            self.prev_grey = cv.CreateImage(self.image_size, 8, 1)
            self.pyramid = cv.CreateImage(self.image_size, 8, 1)
            self.prev_pyramid = cv.CreateImage(self.image_size, 8, 1)
            self.mask = cv.CreateImage(self.image_size, 8, 1)
            self.features = []
            
        """ Create a grey version of the image """
        cv.CvtColor(self.image, self.grey, cv.CV_BGR2GRAY)
        
        """ Equalize the histogram to reduce lighting effects """
        cv.EqualizeHist(self.grey, self.grey)
            
        if self.features != []:
            """ We have feature points, so track and display them """

            """ Calculate the optical flow """
            self.features, status, track_error = cv.CalcOpticalFlowPyrLK(self.prev_grey, self.grey, self.prev_pyramid,
                                                    self.pyramid, self.features, (self.win_size, self.win_size), 3,
                                                    (cv.CV_TERMCRIT_ITER|cv.CV_TERMCRIT_EPS, 20, 0.01), self.flags)

            """ Keep only high status points """
            self.features = [ p for (st,p) in zip(status, self.features) if st]
                                    
        elif self.masks.masks[0].roi:
            """ We do not have features so get the initial features to track.  Note that the
                track_box is initially set to equal the detect_box which is why we can use it 
                for creating a mask below. """
                
            feature_box = None
            
            """ Begin the mask with all black pixels """
            cv.Zero(self.mask)

            """ Get the coordinates and dimensions of the ROI """
            x = self.masks.masks[0].roi.x
            y = self.masks.masks[0].roi.y
            w = self.masks.masks[0].roi.width
            h = self.masks.masks[0].roi.height
            
            if self.auto_face_tracking:
                """ For faces, the detect box tends to extend beyond the actual object so shrink it slightly """
                x = int(0.97 * x)
                y = int(0.97 * y)
                w = int(1 * w)
                h = int(1 * h)
                
                """ Get the center of the track box (type CvRect) so we can create the
                    equivalent CvBox2D (rotated rectangle) required by EllipseBox below. """
                center_x = int(x + w / 2)
                center_y = int(y + h / 2)
                roi_box = ((center_x, center_y), (w, h), 0)
                
                """ Create a white ellipse within the mask to define the ROI.
                    A thickness of -1 causes ellipse to be filled with the chosen color, in this case white. """
                cv.EllipseBox(self.mask, roi_box, cv.CV_RGB(255,255, 255), thickness=-1)
                
            else:
                """ For manually selected regions, just use a rectangle. A thickness of -1 causes ellipse to be filled with
                    the chosen color, in this case white. """
                pt1 = (x, y)
                pt2 = (x + w, y + h)
                cv.Rectangle(self.mask, pt1, pt2, cv.CV_RGB(255,255, 255), thickness=-1)
            
            """ Create the temporary 32-bit scratchpad images """
            eig = cv.CreateImage(self.image_size, 32, 1)
            temp = cv.CreateImage(self.image_size, 32, 1)

            if self.feature_type == 0:
                """ Find points to track using Good Features to Track """
                self.features = cv.GoodFeaturesToTrack(self.grey, eig, temp, self.max_count,
                    self.quality, self.good_feature_distance, mask=self.mask, blockSize=3, useHarris=0, k=0.04)
            
            if self.feature_type == 1 or self.get_surf_also:
                rospy.loginfo("SURF!")
                """ Find keypoints to track using SURF """
                (keypoints, descriptors) = cv.ExtractSURF(self.grey, self.mask, self.SURFMemStorage,
                                                          (0, self.surf_hessian_quality, 3, 1))
                
                """ Set the initial classifier keypoints and descriptors """
                self.classifier_descriptors = descriptors
                self.classifier_keypoints = keypoints
                self.descriptors = descriptors
                self.keypoints = keypoints
                
                if self.feature_type == 1:
                    for keypoint in keypoints:
                        self.features.append(keypoint[0])
            
            if self.auto_min_features:
                """ Since the detect box is larger than the actual face or desired patch, shrink the number a features by 10% """
                self.min_features = int(len(self.features) * 0.9)
                self.abs_min_features = int(0.5 * self.min_features)
        
        """ Swapping the images """
        self.prev_grey, self.grey = self.grey, self.prev_grey
        self.prev_pyramid, self.pyramid = self.pyramid, self.prev_pyramid
        
        """ If we have some features... """
        if len(self.features) > 0:
            """ The FitEllipse2 function below requires us to convert the feature array
                into a CvMat matrix """
            try:
                self.feature_matrix = cv.CreateMat(1, len(self.features), cv.CV_32SC2)
            except:
                pass
                        
            """ Draw the points as green circles and add them to the features matrix """
            i = 0
            for point in self.features:
                detection = Detection()
                if self.feature_type == 0:
                    detection.marker_color = ColorRGBA(0, 255, 0, 0)
                else:
                    detection.marker_color = ColorRGBA(0, 255, 255, 0)     

                detection.label = 'gftt_' + str(point)
                detection.detector = 'GFTT'
                detection.header = Header()
                detection.mask = Mask()
                detection.mask.roi = Rect(point[0], point[1], 1, 1)
                detection.header.stamp = self.image_stamp
                detection.header.frame_id = '/openni_rgb_optical_frame'
                self.detections.detections.append(detection)                       
                try:
                    cv.Set2D(self.feature_matrix, 0, i, (int(point[0]), int(point[1])))
                except:
                    pass
                i = i + 1
    
            """ Draw the best fit ellipse around the feature points """
            if len(self.features) > 6:
                feature_box = cv.FitEllipse2(self.feature_matrix)
            else:
                feature_box = None
            
            """ Publish the ROI for the tracked object """
            try:
                (roi_center, roi_size, roi_angle) = feature_box
            except:
                rospy.loginfo("Patch box has shrunk to zero...")
                feature_box = None
    
#            if feature_box and not self.drag_start and self.is_rect_nonzero(self.track_box):
#                self.ROI = RegionOfInterest()
#                self.ROI.x_offset = int(roi_center[0] - roi_size[0] / 2)
#                self.ROI.y_offset = int(roi_center[1] - roi_size[1] / 2)
#                self.ROI.width = int(roi_size[0])
#                self.ROI.height = int(roi_size[1])
#                
#            self.pubROI.publish(self.ROI)
            
        self.detections.header.stamp = self.image_stamp
        self.detections.header.frame_id = '/openni_rgb_optical_frame' 

        
    def is_rect_nonzero(self, r):
        # First assume a simple CvRect type
        try:
            (_,_,w,h) = r
            return (w > 0) and (h > 0)
        except:
            # Otherwise, assume a CvBox2D type
            ((_,_),(w,h),a) = r
            return (w > 0) and (h > 0) 
    
    def setImage(self, image):
        pass
    
    def setPointCloud(self, point_cloud):
        pass
    
    def loadModels(self):
        pass    

def main(args):
    # Display a help message if appropriate.
    help_message = ""
          
    print help_message
    
    try:
        GoodFeatures("good_features_detector")
        rospy.spin()
    except KeyboardInterrupt:
        print "Shutting down detector node."

if __name__ == '__main__':
    main(sys.argv)